package cn.hughie.hfm;

import java.util.PriorityQueue;

/**
 * @Author yuh
 * @Date Created in 上午6:29 2018/5/23
 * @Description 霍夫曼压缩
 */
public class HFM {

    /**
     * 节点
     */
    private static class Node implements Comparable<Node> {

        Node left;
        Node right;
        char c;
        int freq;


        public Node(Node left, Node right, char c, int freq) {
            this.left = left;
            this.right = right;
            this.c = c;
            this.freq = freq;
        }

        public Node() {
        }

        @Override
        public int compareTo(Node o) {
            return this.freq - o.freq;
        }

        public boolean isLeaf() {
            return left == null && right == null;
        }
    }

    /**
     * 返回频率数组
     *
     * @param s
     * @return
     */
    private int[] calcFreq(String s) {
        int[] freqs = new int[Short.MAX_VALUE + 1];
        char[] chars = s.toCharArray();
        for (char aChar : chars) {
            freqs[aChar]++;
        }
        return freqs;
    }


    /**
     * 构造霍夫曼树
     * 返回根节点
     *
     * @param freq
     * @return
     */
    private Node buildTree(int[] freq) {
        PriorityQueue<Node> queue = new PriorityQueue<>();
        for (int i = 0; i < freq.length; i++) {
            if (freq[i] > 0) {
                queue.add(new Node(null, null, (char) i, freq[i]));
            }
        }

        while (queue.size() > 1) {
            Node left = queue.poll();
            Node right = queue.poll();
            Node parent = new Node(left, right, '\0', left.freq + right.freq);
            queue.add(parent);
        }
        return queue.poll();
    }


    /**
     * 递归构造编译表
     *
     * @param x
     * @param s
     * @param st
     * @return
     */
    private void buildCodeTable(Node x, String s, String[] st) {
        if (x.isLeaf()) {
            st[x.c] = s;
            return;
        }
        buildCodeTable(x.left, s + '0', st);
        buildCodeTable(x.right, s + '1', st);
    }

    /**
     * 前序遍历 就是外围一圈
     * 写一棵树
     *
     * @param x
     */
    private void writeTrie(Node x) {
        if (x.isLeaf()) {
            BinaryStdOut.write(true);
            BinaryStdOut.write(x.c);
            return;
        }
        BinaryStdOut.write(false);
        writeTrie(x.left);
        writeTrie(x.right);
    }

    /**
     * 读一棵树 这是丢失频率的树
     *
     * @return
     */
    private Node readTrie() {
        boolean b = BinaryStdIn.readBoolean();
        if (b) {
            return new Node(null, null, BinaryStdIn.readChar(), 0);
        }
        return new Node(readTrie(), readTrie(), '\0', 0);
    }

    /**
     * 写霍夫曼编码
     *
     * @param n
     * @param st
     */
    private void writeHfmCode(String[] st, char[] input) {
        for (int i = 0; i < input.length; i++) {
            String s = st[input[i]];
            for (int j = 0; j < s.length(); j++) {
                if ('1' == s.charAt(j)) {
                    BinaryStdOut.write(true);
                } else {
                    BinaryStdOut.write(false);
                }
            }
        }
    }

    /**
     * 解压霍夫曼编码
     */
    public void expand() {
        Node node = readTrie();
        int n = BinaryStdIn.readInt();
        for (int i = 0; i < n; i++) {
            if (!node.isLeaf()) {
                if (BinaryStdIn.readBoolean()) {
                    node = node.right;
                } else {
                    node = node.left;
                }
            }
        }
        BinaryStdOut.write(node.c);
    }

    /**
     * 霍夫曼压缩
     */
    public void compress(String s) {
        int[] freqs = calcFreq(s);
        Node node = buildTree(freqs);
        String[] st = new String[Short.MAX_VALUE+1];
        buildCodeTable(node,"",st);
        writeTrie(node);
        BinaryStdOut.write(s.length());
        writeHfmCode(st,s.toCharArray());
    }

    public static void main(String[] args) {
        HFM hfm = new HFM();
        hfm.compress("aaaabc");
        hfm.expand();
    }


}
